package com.example.demo.util;
/**
 * Copyright: Copyright (c) 2021 umbpay. All Rights Reserved.
 * <p>
 * 功能描述：
 *
 * @author umbpay
 * Create Date: 2022/6/15  19:17
 * @version v1.0.0
 */

import cfca.sadk.algorithm.common.*;
import cfca.sadk.algorithm.sm2.SM4Engine;
import cfca.sadk.algorithm.util.BigFileCipherUtil;
import cfca.sadk.algorithm.util.SM2AndItsCloseSymAlgUtil;
import cfca.sadk.asn1.parser.ASN1Node;
import cfca.sadk.asn1.parser.ASN1Parser;
import cfca.sadk.asn1.parser.EnvelopFileParser;
import cfca.sadk.envelope.sm2.SM2EncryptedContentInfo;
import cfca.sadk.envelope.sm2.SM2EncryptedInputStream;
import cfca.sadk.envelope.sm2.SM2EnvelopedData;
import cfca.sadk.lib.crypto.Session;
import cfca.sadk.lib.crypto.jni.JNISoftLib;
import cfca.sadk.org.bouncycastle.asn1.*;
import cfca.sadk.org.bouncycastle.asn1.cms.*;
import cfca.sadk.org.bouncycastle.asn1.x500.X500Name;
import cfca.sadk.org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import cfca.sadk.org.bouncycastle.asn1.x509.SubjectKeyIdentifier;
import cfca.sadk.org.bouncycastle.cms.CMSEnvelopedData;
import cfca.sadk.util.Base64;
import cfca.sadk.x509.certificate.X509Cert;

import java.io.File;
import java.io.FileOutputStream;
import java.math.BigInteger;
import java.security.PrivateKey;

public class SM2EnvelopeUtil {
    private static byte[] IV_16 = new byte[]{50, 51, 52, 53, 54, 55, 56, 57, 56, 55, 54, 53, 52, 51, 50, 49};

    public SM2EnvelopeUtil() {
    }

    public static final void envelopeFile(String sourceFilePath, String encryptedFilePath, String symmetricAlgorithm, X509Cert[] receiverCerts) throws Exception {
        envelopeFile_Session(sourceFilePath, encryptedFilePath, symmetricAlgorithm, receiverCerts, (Session)null);
    }

    public static final void envelopeFile(String sourceFilePath, String encryptedFilePath, String symmetricAlgorithm, X509Cert[] receiverCerts, Session session) throws Exception {
        envelopeFile_Session(sourceFilePath, encryptedFilePath, symmetricAlgorithm, receiverCerts, session);
    }

    private static final void envelopeFile_Session(String sourceFilePath, String encryptedFilePath, String symmetricAlgorithm, X509Cert[] receiverCerts, Session session) throws Exception {
        byte[] key = SM2AndItsCloseSymAlgUtil.generateSecretKey();
        IV_16 = SM2AndItsCloseSymAlgUtil.generateIV();
        ASN1EncodableVector recipientInfos = new ASN1EncodableVector();

        for(int i = 0; i < receiverCerts.length; ++i) {
            recipientInfos.add(toRecipientInfo(receiverCerts[i], key, session));
            recipientInfos.add(toRecipientInfoOfIssuerAndSerialNumber(receiverCerts[i], key, session));
        }

        Mechanism contentEncryptionAlg;
        if (symmetricAlgorithm.indexOf("CBC") != -1) {
            CBCParam cbc = new CBCParam(IV_16);
            contentEncryptionAlg = new Mechanism(symmetricAlgorithm, cbc);
        } else {
            contentEncryptionAlg = new Mechanism(symmetricAlgorithm);
        }

        ASN1ObjectIdentifier tOID = (ASN1ObjectIdentifier)PKCS7EnvelopedData.MECH_OID.get(symmetricAlgorithm);
        AlgorithmIdentifier algId = getAlgorithmIdentifier(contentEncryptionAlg, tOID);
        boolean useJNIFlag = false;
        if (session != null && session instanceof JNISoftLib) {
            useJNIFlag = true;
        }

        SM2EncryptedInputStream encryptStream = new SM2EncryptedInputStream(useJNIFlag, new File(sourceFilePath), key, contentEncryptionAlg);
        SM2EncryptedContentInfo sm2Eci = new SM2EncryptedContentInfo(PKCSObjectIdentifiers.sm2Data, algId, encryptStream);
        SM2EnvelopedData enData = new SM2EnvelopedData((OriginatorInfo)null, new DERSet(recipientInfos), sm2Eci, (ASN1Set)null);
        ContentInfo contentInfo = new ContentInfo(PKCSObjectIdentifiers.sm2EnvelopedData, enData);
        File f = new File(encryptedFilePath);
        if (!f.exists()) {
            f.createNewFile();
        }

        FileOutputStream fos = new FileOutputStream(f);
        DEROutputStream dos = new DEROutputStream(fos);
        dos.writeObject(contentInfo);
        dos.close();
    }

    public static final void openEnvelopedFile(String encryptedFilePath, String plainTextFilePath, PrivateKey privateKey, X509Cert recipientCert) throws Exception {
        EnvelopFileParser parser = new EnvelopFileParser(new File(encryptedFilePath));
        parser.parser();
        ASN1Node receiver_node = parser.getReceiver_node();
        ASN1Node encrypted_node = parser.getEncrypted_node();
        openEnvelopFile_ASN1Node(receiver_node, encrypted_node, privateKey, recipientCert, plainTextFilePath);
    }

    private static boolean hasRecipent_file(KeyTransRecipientInfo inf, byte[] subjectPubKeyID, X500Name recipientIssuer, BigInteger recipientSN) {
        RecipientIdentifier id = inf.getRecipientIdentifier();
        DEROctetString oct = new DEROctetString(subjectPubKeyID);
        IssuerAndSerialNumber issu = new IssuerAndSerialNumber(recipientIssuer, recipientSN);
        return id.getId().toASN1Primitive().asn1Equals(oct) || id.getId().toASN1Primitive().asn1Equals(issu.toASN1Primitive());
    }

    private static void openEnvelopFile_ASN1Node(ASN1Node receiver_node, ASN1Node encrypted_node, PrivateKey privateKey, X509Cert recipientCert, String plainTextFilePath) throws Exception {
        FileOutputStream fos = null;

        try {
            X500Name recipientIssuer = recipientCert.getIssuerX500Name();
            BigInteger recipientSN = recipientCert.getSerialNumber();
            byte[] subjectPubKeyID = recipientCert.getSubjectKeyIdentifier().getKeyIdentifier();
            ASN1Set receivers = ASN1Set.getInstance(receiver_node.getData());
            int len = receivers.size();
            ASN1OctetString encryptKey = null;
            AlgorithmIdentifier algId = null;

            for(int i = 0; i < len; ++i) {
                RecipientInfo recip = RecipientInfo.getInstance(receivers.getObjectAt(i));
                if (recip.getInfo() instanceof KeyTransRecipientInfo) {
                    KeyTransRecipientInfo inf = KeyTransRecipientInfo.getInstance(recip.getInfo());
                    if (hasRecipent_file(inf, subjectPubKeyID, recipientIssuer, recipientSN)) {
                        encryptKey = inf.getEncryptedKey();
                        algId = inf.getKeyEncryptionAlgorithm();
                        break;
                    }
                }
            }

            if (encryptKey == null || algId == null) {
                throw new Exception("can not find the receiver!!!");
            }

            byte[] symmetricKey = SM2AndItsCloseSymAlgUtil.sm2Encrypt(false, privateKey, encryptKey.getOctets());
            ASN1Node symmetric_encrypted = (ASN1Node)encrypted_node.childNodes.get(1);
            byte[] symmetricAlg_byte = symmetric_encrypted.getData();
            AlgorithmIdentifier symmetricAlgId = AlgorithmIdentifier.getInstance(ASN1Sequence.getInstance(symmetricAlg_byte));
            String encryptionAlgStr = (String)PKCS7EnvelopedData.OID_MECH.get(symmetricAlgId.getAlgorithm());
            Mechanism mechanism = null;
            if (encryptionAlgStr.indexOf("CBC") != -1) {
                DEROctetString doct = (DEROctetString)symmetricAlgId.getParameters();
                CBCParam cbcParam = new CBCParam(doct.getOctets());
                if (encryptionAlgStr.equals("SM4/CBC/PKCS7Padding")) {
                    mechanism = new Mechanism("SM4/CBC/PKCS7Padding", cbcParam);
                }
            } else {
                if (encryptionAlgStr.indexOf("ECB") == -1) {
                    throw new PKIException(PKIException.UNSUPPORT_ENCRYPT_ALG_SIGNANDENVELOP_ERR, PKIException.UNSUPPORT_ENCRYPT_ALG_SIGNANDENVELOP_ERR_DES + "Algorithm is:" + encryptionAlgStr);
                }

                mechanism = new Mechanism("SM4/ECB/PKCS7Padding");
            }

            if (mechanism == null) {
                throw new PKIException(PKIException.UNSUPPORT_ENCRYPT_ALG_SIGNANDENVELOP_ERR, PKIException.UNSUPPORT_ENCRYPT_ALG_SIGNANDENVELOP_ERR_DES + "Algorithm is:" + encryptionAlgStr);
            }

            File f = new File(plainTextFilePath);
            if (!f.exists()) {
                f.createNewFile();
            }

            fos = new FileOutputStream(f);
            ASN1Node file_encrypted = (ASN1Node)encrypted_node.childNodes.get(2);
            if (file_encrypted.childNodes.size() == 1) {
                file_encrypted = (ASN1Node)file_encrypted.childNodes.get(0);
            }

            BigFileCipherUtil.bigFileBlockDecrypt(symmetricKey, new SM4Engine(), (CBCParam)mechanism.getParam(), file_encrypted, fos);
        } finally {
            if (fos != null) {
                fos.close();
            }

        }

    }

    public static final byte[] envelopeMessage(byte[] sourceData, String symmetricAlgorithm, X509Cert[] receiverCerts) throws Exception {
        byte[] none64 = envelopMessage_None64(sourceData, symmetricAlgorithm, receiverCerts, (Session)null);
        return Base64.encode(none64);
    }

    public static final byte[] envelopeMessage(byte[] sourceData, String symmetricAlgorithm, X509Cert[] receiverCerts, Session session) throws Exception {
        byte[] none64 = envelopMessage_None64(sourceData, symmetricAlgorithm, receiverCerts, session);
        return Base64.encode(none64);
    }

    private static byte[] envelopMessage_None64(byte[] sourceData, String symmetricAlgorithm, X509Cert[] receiverCerts, Session session) throws Exception {
        byte[] key = SM2AndItsCloseSymAlgUtil.generateSecretKey();
        IV_16 = SM2AndItsCloseSymAlgUtil.generateIV();
        ASN1EncodableVector recipientInfos = new ASN1EncodableVector();

        for(int i = 0; i < receiverCerts.length; ++i) {
            //recipientInfos.add(toRecipientInfo(receiverCerts[i], key, session));
            recipientInfos.add(toRecipientInfoOfIssuerAndSerialNumber(receiverCerts[i], key, session));
        }

        Mechanism contentEncryptionAlg;
        if (symmetricAlgorithm.indexOf("CBC") != -1) {
            CBCParam cbc = new CBCParam(IV_16);
            contentEncryptionAlg = new Mechanism(symmetricAlgorithm, cbc);
        } else {
            contentEncryptionAlg = new Mechanism(symmetricAlgorithm);
        }

        boolean useJNI = false;
        if (session != null && session instanceof JNISoftLib) {
            useJNI = true;
        }

        byte[] encryptedData = SM2AndItsCloseSymAlgUtil.crypto(useJNI, true, key, sourceData, contentEncryptionAlg);
        ASN1OctetString encryptedOctet = new BEROctetString(encryptedData);
        ASN1ObjectIdentifier tOID = (ASN1ObjectIdentifier)PKCS7EnvelopedData.MECH_OID.get(symmetricAlgorithm);
        AlgorithmIdentifier algId = getAlgorithmIdentifier(contentEncryptionAlg, tOID);
        ASN1Encodable parameters = algId.getParameters();
        EncryptedContentInfo encryptedContentInfo = new EncryptedContentInfo(PKCSObjectIdentifiers.sm2Data, algId, encryptedOctet);
        EnvelopedData envData = new EnvelopedData((OriginatorInfo)null, new DERSet(recipientInfos), encryptedContentInfo, ASN1Set.getInstance((Object)null));
        ContentInfo contentInfo = new ContentInfo(PKCSObjectIdentifiers.sm2EnvelopedData, envData);
        return ASN1Parser.parseDERObj2Bytes((new CMSEnvelopedData(contentInfo)).toASN1Structure());
    }

    private static AlgorithmIdentifier getAlgorithmIdentifier(Mechanism contentEncryptionAlg, ASN1ObjectIdentifier tOID) throws PKIException {
        AlgorithmIdentifier algorithmIdentifier = new AlgorithmIdentifier(tOID);
        if (contentEncryptionAlg.getMechanismType().toUpperCase().indexOf("CBC") != -1) {
            Object param = contentEncryptionAlg.getParam();
            if (param == null) {
                throw new PKIException(PKIException.NULL_P7_ENVELOP_CBC_ERR, PKIException.NULL_P7_ENVELOP_CBC_ERR_DES);
            } else {
                CBCParam cbcParam = (CBCParam)contentEncryptionAlg.getParam();
                DEROctetString doct = new DEROctetString(cbcParam.getIv());
                algorithmIdentifier = new AlgorithmIdentifier(tOID, doct);
                return algorithmIdentifier;
            }
        } else {
            return algorithmIdentifier;
        }
    }

    private static RecipientInfo toRecipientInfo(X509Cert cert, byte[] symmetricKey, Session session) throws Exception {
        byte[] encryptedKey = null;
        if (session != null && session instanceof JNISoftLib) {
            encryptedKey = SM2AndItsCloseSymAlgUtil.sm2EncryptByJNI(true, cert.getPublicKey(), symmetricKey);
        } else {
            encryptedKey = SM2AndItsCloseSymAlgUtil.sm2Encrypt(true, cert.getPublicKey(), symmetricKey);
        }

        ASN1OctetString encKey = new DEROctetString(encryptedKey);
        SubjectKeyIdentifier sid = cert.getSubjectKeyIdentifier();
        if (sid == null) {
            throw new Exception("the cert has no extension data with SubjectKeyIdentifier,can not create envelope data");
        } else {
            AlgorithmIdentifier keyEncAlg = new AlgorithmIdentifier(PKCSObjectIdentifiers.SM2_pubKey_encrypt, DERNull.INSTANCE);

            RecipientIdentifier recipientIdentifier = new RecipientIdentifier(new DEROctetString(sid.getKeyIdentifier()));


            KeyTransRecipientInfo ktr = new KeyTransRecipientInfo(recipientIdentifier, keyEncAlg, encKey);

            return new RecipientInfo(ktr);
        }
    }

    private static RecipientInfo toRecipientInfoOfIssuerAndSerialNumber(X509Cert cert, byte[] symmetricKey, Session session) throws Exception {

        byte[] encryptedKey = null;
        if (session != null && session instanceof JNISoftLib) {
            encryptedKey = SM2AndItsCloseSymAlgUtil.sm2EncryptByJNI(true, cert.getPublicKey(), symmetricKey);
        } else {
            encryptedKey = SM2AndItsCloseSymAlgUtil.sm2Encrypt(true, cert.getPublicKey(), symmetricKey);
        }

        ASN1OctetString encKey = new DEROctetString(encryptedKey);
        //SubjectKeyIdentifier sid = cert.getSubjectKeyIdentifier();
//        if (sid == null) {
//            throw new Exception("the cert has no extension data with SubjectKeyIdentifier,can not create envelope data");
//        } else {
        X500Name recipientIssuer = cert.getIssuerX500Name();
        BigInteger recipientSN = cert.getSerialNumber();
        IssuerAndSerialNumber issu = new IssuerAndSerialNumber(recipientIssuer, recipientSN);

        AlgorithmIdentifier keyEncAlg = new AlgorithmIdentifier(PKCSObjectIdentifiers.SM2_pubKey_encrypt, DERNull.INSTANCE);

        //RecipientIdentifier recipientIdentifier = new RecipientIdentifier(new DEROctetString(sid));
        KeyTransRecipientInfo ktr = new KeyTransRecipientInfo(RecipientIdentifier.getInstance(issu), keyEncAlg, encKey);

        return new RecipientInfo(ktr);
        //}
    }
}